package io.gitee.luoyuan800.tool.llm.agent;

import io.gitee.luoyuan800.tool.llm.bean.AbstractMessage;
import io.gitee.luoyuan800.tool.llm.bean.ImageMessage;
import io.gitee.luoyuan800.tool.llm.bean.Message;
import io.gitee.luoyuan800.tool.llm.bean.completion.ChatCompletion;
import io.gitee.luoyuan800.tool.llm.bean.completion.CommonChatCompletion;
import io.gitee.luoyuan800.tool.llm.bean.completion.DeltaChatCompletion;
import io.gitee.luoyuan800.tool.llm.bean.config.AiAgentConfig;
import io.gitee.luoyuan800.tool.llm.bean.request.ChatRequest;
import io.gitee.luoyuan800.tool.llm.web.RestfulClient;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Stream;

@Data
@RequiredArgsConstructor
@Slf4j
public class CommonAgent implements AiAgent {
    private final List<AbstractMessage<?>> messages = new ArrayList<>();
    private final AiAgentConfig aiAgentConfig;
    private final RestfulClient restTemplate;
    private final String name;
    private final String model;
    private final AtomicLong tokenUsage = new AtomicLong(0);

    public CommonAgent(AiAgentConfig aiAgentConfig, RestfulClient restTemplate, String name) {
        this.aiAgentConfig = aiAgentConfig;
        this.restTemplate = restTemplate;
        this.name = name;
        this.model = aiAgentConfig.getModel();
        if (StringUtils.isNoneBlank(aiAgentConfig.getSystem())) {
            addMessage(ChatRequest.MessageType.SYSTEM, aiAgentConfig.getSystem(), null);
        }
    }

    public void applySelectResponse(Message msg) {
        if (msg != null && aiAgentConfig.isHistory()) {
            this.messages.add(msg);
        }
    }

    public String chatWithImage(String msg, String imageUrl, boolean enableThinking) {
        CommonChatCompletion completion = fullChat(msg, 1, imageUrl, enableThinking);
        if (completion != null && completion.getChoices() != null && !completion.getChoices().isEmpty()) {
            return buildMessageString(true, completion.getChoices().get(0).getMessage());
        }
        return "";
    }

    @Override
    public String chatWithImage(String msg, String image) {
        return chatWithImage(msg, image, aiAgentConfig.isEnableThink());
    }

    public CommonChatCompletion fullChat(String msg, int resultCount, String image, boolean enableThinking) {
        log.info("Sending message to {}({}).", name, model);
        log.debug("Message: {}", msg);

        ChatRequest request = buildRequest(resultCount);
        request.setStream(false);

        List<AbstractMessage<?>> messageSend = addMessage(ChatRequest.MessageType.USER, msg, image);
        request.setMessages(messageSend);

        CommonChatCompletion response = httpRequest(request);

        tokenTrack(response);
        log.info("Got response from {}({}).", name, model);
        return response;
    }

    void tokenTrack(ChatCompletion<?> response) {
        if (response != null && response.getUsage()!=null) {
            log.info("Agent {}({}) already used {} tokens", name, model, tokenUsage.addAndGet(response.getUsage().getTotalTokens()));
        }
    }

    CommonChatCompletion httpRequest(ChatRequest request) {
        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", "Bearer " + aiAgentConfig.getToken());
        headers.put("date", ZonedDateTime.now().toString());
        headers.put("x-bce-date", ZonedDateTime.now().toString());
        return restTemplate.post(
                aiAgentConfig.getUrl(),
                headers, request,
                CommonChatCompletion.class
        );
    }

    ChatRequest buildRequest(int resultCount) {
        ChatRequest request = new ChatRequest();
        setRequestParams(resultCount, request);
        return request;
    }

    void setRequestParams(int resultCount, ChatRequest request) {
        request.setModel(model);
        request.setStream(aiAgentConfig.isStream());
        request.setTemperature(aiAgentConfig.getTemperature());
        request.setStop(aiAgentConfig.getStop());
        request.setMaxToken(aiAgentConfig.getMaxToken());
        request.setMaxCompletionTokens(aiAgentConfig.getMaxToken());
        request.setFrequencyPenalty(aiAgentConfig.getFrequencyPenalty());
        request.setPresencePenalty(aiAgentConfig.getPresencePenalty());
        request.setN(resultCount);
        request.setTopP(aiAgentConfig.getTopP());
        request.setPenaltyScore(aiAgentConfig.getPenaltyScore());
        request.setResponseFormat(aiAgentConfig.getResponseFormat());
    }

    public String chat(String msg) {
        return chat(msg, aiAgentConfig.isEnableThink());
    }

    public String chat(String msg, boolean showThink) {
        Message response = chatReturnObject(msg, showThink);
        return buildMessageString(showThink, response);
    }

    String buildMessageString(boolean showThink, Message response) {
        if (response != null) {
            return showThink ? response.getFullContent() : response.getContent();
        }
        return "";
    }

    private Message chatReturnObject(String msg, boolean enableThinking) {
        List<Message> messageList = chatMultiple(msg, 1, enableThinking);
        Message message = messageList
                .stream().findFirst()
                .orElse(null);
        applySelectResponse(message);
        return message;
    }

    public List<Message> chatMultiple(String msg, int times, boolean enableThinking) {
        CommonChatCompletion commonChatCompletion = fullChat(msg, times, null, enableThinking);
        return Optional.ofNullable(commonChatCompletion.getChoices()).orElse(List.of())
                .stream()
                .map(CommonChatCompletion.NormalChoice::getMessage)
                .toList();
    }

    List<AbstractMessage<?>> addMessage(ChatRequest.MessageType messageType, String msg, String image) {
        AbstractMessage<?> message;
        if (StringUtils.isNotBlank(image)) {
            message = new ImageMessage(msg, image);
        } else {
            Message contentMessage = new Message();
            contentMessage.setContent(msg);
            message = contentMessage;
        }
        message.setRole(messageType.getId());
        List<AbstractMessage<?>> msgs = new ArrayList<>(messages);
        msgs.add(message);
        if (aiAgentConfig.isHistory()) {
            messages.add(message);
        }
        return msgs;
    }

    public Stream<Message> streamChat(String msg) {
        ChatRequest request = buildRequest(1);
        setRequestParams(1, request);
        request.setStream(true);

        addMessage(ChatRequest.MessageType.USER, msg, null);
        List<AbstractMessage<?>> newMessage = new ArrayList<>(messages);
        request.setMessages(newMessage);

        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", "Bearer " + aiAgentConfig.getToken());
        headers.put("date", ZonedDateTime.now().toString());
        headers.put("x-bce-date", ZonedDateTime.now().toString());
        return restTemplate.postStream(
                        aiAgentConfig.getUrl(),
                        headers, request,
                        DeltaChatCompletion.class
                ).map(delta -> {
                    tokenTrack(delta);
                    return delta.getChoices();
                }).filter(Objects::nonNull)
                .flatMap(List::stream)
                .map(DeltaChatCompletion.DeltaChoice::getDelta)
                .filter(Objects::nonNull);
    }

}
